from lcb_runner_cq.lm_styles import LMStyle


# def extract_code(model_output: str, lmstyle: LMStyle):
#     outputlines = model_output.split("\n")
#     if lmstyle == LMStyle.CodeLLaMaInstruct:
#         indexlines = [i for i, line in enumerate(outputlines) if "PYTHON]" in line]
#         if len(indexlines) < 2:
#             indexlines = [i for i, line in enumerate(outputlines) if "```" in line]
#     elif lmstyle in [
#         LMStyle.CodeLLaMaBase,
#         LMStyle.DeepSeekBase,
#         LMStyle.StarCoder2Base,
#         LMStyle.StableCodeBase,
#         LMStyle.CodeQwenBase,
#     ]:
#         return model_output.strip()
#     else:
#         indexlines = [i for i, line in enumerate(outputlines) if "```" in line]
#     if len(indexlines) < 2:
#         return ""
#     return "\n".join(outputlines[indexlines[0] + 1 : indexlines[1]])


def extract_code(model_output: str, lmstyle: LMStyle):
    outputlines = model_output.split("\n")
    indexlines = [i for i, line in enumerate(outputlines) if "```" in line]
    if len(indexlines) < 2:
        return ""
    return "\n".join(outputlines[indexlines[-2] + 1 : indexlines[-1]])


def extract_test_output_code(model_output: str, lmstyle: LMStyle = None):
    outputlines = model_output.split("\n")
    # find the last line startwith assert...
    indexlines = [i for i, line in enumerate(outputlines) if line.startswith("assert")]
    if indexlines:
        return outputlines[indexlines[-1]]
    if lmstyle and lmstyle == LMStyle.CodeLLaMaInstruct:
        indexlines = [i for i, line in enumerate(outputlines) if "PYTHON]" in line]
    else:
        # first try to extract ```python if not then try ```
        indexlines = [i for i, line in enumerate(outputlines) if "```python" in line or "```Python" in line]
        if indexlines:
            start_index = indexlines[0]
        else:
            start_index = None
        indexlines = [i for i, line in enumerate(outputlines) if "```" in line]
        if start_index is not None:
            indexlines = [i for i in indexlines if i > start_index]
            indexlines = [start_index] + indexlines

    if len(indexlines) < 2:
        return ""
    return "\n".join(outputlines[indexlines[0] + 1 : indexlines[1]])
